{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "torch_io tutorial for transf3d seg.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RFGq4NTRNfyJ"
      },
      "source": [
        "## self-attention-cv : illustration of a training process with subvolume sampling for 3d segmentation\n",
        "\n",
        "The dataset can be found here: https://iseg2019.web.unc.edu/ . i uploaded it and mounted from my gdrive"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "clse_QTnggr-",
        "outputId": "146e13ec-89ef-40af-f6b4-a6c4e92bbb5f"
      },
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/gdrive')\n",
        "import zipfile\n",
        "root_path = '/gdrive/My Drive/DATASETS/iSeg-2019-Training.zip' \n",
        "!echo \"Download and extracting folders...\"\n",
        "zip_ref = zipfile.ZipFile(root_path, 'r')\n",
        "zip_ref.extractall(\"./\")\n",
        "zip_ref.close()\n",
        "!echo \"Finished\"\n",
        "!pip install torchio\n",
        "!pip install self-attention-cv"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Mounted at /gdrive\n",
            "Download and extracting folders...\n",
            "Finished\n",
            "Collecting torchio\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/3d/33/94812ae74a2815fdd5bf7c4e26be75086ebc770309c569380e6f7cc4ad60/torchio-0.18.29-py2.py3-none-any.whl (140kB)\n",
            "\u001b[K     |████████████████████████████████| 143kB 18.9MB/s \n",
            "\u001b[?25hRequirement already satisfied: nibabel in /usr/local/lib/python3.7/dist-packages (from torchio) (3.0.2)\n",
            "Requirement already satisfied: Click in /usr/local/lib/python3.7/dist-packages (from torchio) (7.1.2)\n",
            "Collecting Deprecated\n",
            "  Downloading https://files.pythonhosted.org/packages/d4/56/7d4774533d2c119e1873993d34d313c9c9efc88c5e4ab7e33bdf915ad98c/Deprecated-1.2.11-py2.py3-none-any.whl\n",
            "Requirement already satisfied: tqdm in /usr/local/lib/python3.7/dist-packages (from torchio) (4.41.1)\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.7/dist-packages (from torchio) (1.19.5)\n",
            "Requirement already satisfied: torch>=1.1 in /usr/local/lib/python3.7/dist-packages (from torchio) (1.8.0+cu101)\n",
            "Collecting SimpleITK<2\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/4a/ee/638b6bae2db10e5ef4ca94c95bb29ec25aa37a9d721b47f91077d7e985e0/SimpleITK-1.2.4-cp37-cp37m-manylinux1_x86_64.whl (42.5MB)\n",
            "\u001b[K     |████████████████████████████████| 42.5MB 65kB/s \n",
            "\u001b[?25hRequirement already satisfied: humanize in /usr/local/lib/python3.7/dist-packages (from torchio) (0.5.1)\n",
            "Requirement already satisfied: scipy in /usr/local/lib/python3.7/dist-packages (from torchio) (1.4.1)\n",
            "Requirement already satisfied: wrapt<2,>=1.10 in /usr/local/lib/python3.7/dist-packages (from Deprecated->torchio) (1.12.1)\n",
            "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.7/dist-packages (from torch>=1.1->torchio) (3.7.4.3)\n",
            "Installing collected packages: Deprecated, SimpleITK, torchio\n",
            "Successfully installed Deprecated-1.2.11 SimpleITK-1.2.4 torchio-0.18.29\n",
            "Collecting self-attention-cv\n",
            "  Downloading https://files.pythonhosted.org/packages/69/5b/4163230c657f80a5f4123af111a410257df5d0c3d26027bd150c40b75fec/self_attention_cv-1.1.0-py3-none-any.whl\n",
            "Collecting einops>=0.3\n",
            "  Downloading https://files.pythonhosted.org/packages/5d/a0/9935e030634bf60ecd572c775f64ace82ceddf2f504a5fd3902438f07090/einops-0.3.0-py2.py3-none-any.whl\n",
            "Requirement already satisfied: numpy>=1.19 in /usr/local/lib/python3.7/dist-packages (from self-attention-cv) (1.19.5)\n",
            "Collecting pytest>=6.2\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/a1/cf/7f67585bd2fc0359ec482cf3c430bce3ef6d3f40bc468137225a733e3069/pytest-6.2.2-py3-none-any.whl (280kB)\n",
            "\u001b[K     |████████████████████████████████| 286kB 14.2MB/s \n",
            "\u001b[?25hRequirement already satisfied: torch>=1.7 in /usr/local/lib/python3.7/dist-packages (from self-attention-cv) (1.8.0+cu101)\n",
            "Requirement already satisfied: torchvision>=0.8 in /usr/local/lib/python3.7/dist-packages (from self-attention-cv) (0.9.0+cu101)\n",
            "Requirement already satisfied: importlib-metadata>=0.12; python_version < \"3.8\" in /usr/local/lib/python3.7/dist-packages (from pytest>=6.2->self-attention-cv) (3.7.0)\n",
            "Requirement already satisfied: attrs>=19.2.0 in /usr/local/lib/python3.7/dist-packages (from pytest>=6.2->self-attention-cv) (20.3.0)\n",
            "Requirement already satisfied: packaging in /usr/local/lib/python3.7/dist-packages (from pytest>=6.2->self-attention-cv) (20.9)\n",
            "Requirement already satisfied: toml in /usr/local/lib/python3.7/dist-packages (from pytest>=6.2->self-attention-cv) (0.10.2)\n",
            "Requirement already satisfied: py>=1.8.2 in /usr/local/lib/python3.7/dist-packages (from pytest>=6.2->self-attention-cv) (1.10.0)\n",
            "Collecting pluggy<1.0.0a1,>=0.12\n",
            "  Downloading https://files.pythonhosted.org/packages/a0/28/85c7aa31b80d150b772fbe4a229487bc6644da9ccb7e427dd8cc60cb8a62/pluggy-0.13.1-py2.py3-none-any.whl\n",
            "Requirement already satisfied: iniconfig in /usr/local/lib/python3.7/dist-packages (from pytest>=6.2->self-attention-cv) (1.1.1)\n",
            "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.7/dist-packages (from torch>=1.7->self-attention-cv) (3.7.4.3)\n",
            "Requirement already satisfied: pillow>=4.1.1 in /usr/local/lib/python3.7/dist-packages (from torchvision>=0.8->self-attention-cv) (7.0.0)\n",
            "Requirement already satisfied: zipp>=0.5 in /usr/local/lib/python3.7/dist-packages (from importlib-metadata>=0.12; python_version < \"3.8\"->pytest>=6.2->self-attention-cv) (3.4.1)\n",
            "Requirement already satisfied: pyparsing>=2.0.2 in /usr/local/lib/python3.7/dist-packages (from packaging->pytest>=6.2->self-attention-cv) (2.4.7)\n",
            "\u001b[31mERROR: datascience 0.10.6 has requirement folium==0.2.1, but you'll have folium 0.8.3 which is incompatible.\u001b[0m\n",
            "Installing collected packages: einops, pluggy, pytest, self-attention-cv\n",
            "  Found existing installation: pluggy 0.7.1\n",
            "    Uninstalling pluggy-0.7.1:\n",
            "      Successfully uninstalled pluggy-0.7.1\n",
            "  Found existing installation: pytest 3.6.4\n",
            "    Uninstalling pytest-3.6.4:\n",
            "      Successfully uninstalled pytest-3.6.4\n",
            "Successfully installed einops-0.3.0 pluggy-0.13.1 pytest-6.2.2 self-attention-cv-1.1.0\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ziryVMj3Z0Za"
      },
      "source": [
        "## Training example"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Q_AgBbfvh6DV"
      },
      "source": [
        "import glob\n",
        "import torchio as tio\n",
        "import torch\n",
        "from torch.utils.data import DataLoader\n",
        "\n",
        "paths_t1 = sorted(glob.glob('./iSeg-2019-Training/*T1.img'))\n",
        "paths_t2 = sorted(glob.glob('./iSeg-2019-Training/*T2.img'))\n",
        "paths_seg = sorted(glob.glob('./iSeg-2019-Training/*label.img'))\n",
        "assert len(paths_t1) == len(paths_t2) == len(paths_seg)\n",
        "\n",
        "subject_list = []\n",
        "for pat in zip(paths_t1, paths_t2, paths_seg):\n",
        "  path_t1, path_t2, path_seg = pat\n",
        "  subject = tio.Subject(t1=tio.ScalarImage(path_t1,),\n",
        "              t2=tio.ScalarImage(path_t2,),\n",
        "              label=tio.LabelMap(path_seg)) \n",
        "  subject_list.append(subject)\n",
        "\n",
        "\n",
        "transforms = [tio.RescaleIntensity((0, 1)),tio.RandomAffine() ]\n",
        "transform = tio.Compose(transforms)\n",
        "\n",
        "subjects_dataset = tio.SubjectsDataset(subject_list, transform=transform)\n",
        "\n",
        "patch_size = 24\n",
        "queue_length = 300\n",
        "samples_per_volume = 50\n",
        "sampler = tio.data.UniformSampler(patch_size)\n",
        "\n",
        "patches_queue = tio.Queue(\n",
        "subjects_dataset,\n",
        "queue_length,\n",
        "samples_per_volume,sampler, num_workers=1)\n",
        "\n",
        "patches_loader = DataLoader(patches_queue, batch_size=16)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "9vTzIMmWPSj9",
        "outputId": "a5bd1f43-2075-4a5e-c69d-6d93f866522d"
      },
      "source": [
        "from self_attention_cv.Transformer3Dsegmentation import Transformer3dSeg\n",
        "\n",
        "def crop_target(img, target_size):\n",
        "  dim = img.shape[-1]\n",
        "  center = dim//2\n",
        "  start_dim = center - (target_size//2) - 1\n",
        "  end_dim = center + (target_size//2)\n",
        "  return img[:,0,start_dim:end_dim,start_dim:end_dim,start_dim:end_dim].long()\n",
        "\n",
        "target_size = 3 # as in the paper \n",
        "patch_dim = 8\n",
        "num_epochs = 50\n",
        "num_classes = 4\n",
        "model = Transformer3dSeg(subvol_dim=patch_size, patch_dim=patch_dim,\n",
        "                         in_channels=2, blocks=2, num_classes=num_classes).cuda()\n",
        "\n",
        "criterion = torch.nn.CrossEntropyLoss()\n",
        "optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)\n",
        "print(len(patches_loader))\n",
        "for epoch_index in range(num_epochs):\n",
        "  epoch_loss = 0\n",
        "  for c,patches_batch in enumerate(patches_loader):\n",
        "    optimizer.zero_grad()\n",
        "    \n",
        "    input_t1 = patches_batch['t1'][tio.DATA]  \n",
        "    input_t2 = patches_batch['t2'][tio.DATA]\n",
        "\n",
        "    input_tensor = torch.cat([input_t1, input_t2], dim=1).cuda()\n",
        "    \n",
        "    \n",
        "    logits = model(input_tensor) # 8x8x8 the 3d transformer-based approach\n",
        "\n",
        "    # for the 3d transformer-based approach the target must be cropped again to the desired size\n",
        "    targets = patches_batch['label'][tio.DATA]  \n",
        "    \n",
        "    cropped_target = crop_target(targets, target_size).cuda()\n",
        "\n",
        "    loss = criterion(logits, cropped_target)\n",
        "    loss.backward()\n",
        "    optimizer.step()\n",
        "    epoch_loss = epoch_loss+loss.cpu().item()\n",
        "\n",
        "  print(f'epoch {epoch_index} loss {epoch_loss/c}')\n",
        "    \n",
        "\n"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "32\n",
            "epoch 0 loss 0.8919196542232267\n",
            "epoch 1 loss 0.6648283805097303\n",
            "epoch 2 loss 0.6422034237653979\n",
            "epoch 3 loss 0.5969387196725414\n",
            "epoch 4 loss 0.5559082502318967\n",
            "epoch 5 loss 0.49828739656556037\n",
            "epoch 6 loss 0.48543436681070634\n",
            "epoch 7 loss 0.3903121284900173\n",
            "epoch 8 loss 0.38039007951175013\n",
            "epoch 9 loss 0.2883441626064239\n",
            "epoch 10 loss 0.35982790421093663\n",
            "epoch 11 loss 0.2505160081650942\n",
            "epoch 12 loss 0.2158138483402229\n",
            "epoch 13 loss 0.20691758676642372\n",
            "epoch 14 loss 0.20189064626972522\n",
            "epoch 15 loss 0.24909109192629975\n",
            "epoch 16 loss 0.18076440347959438\n",
            "epoch 17 loss 0.23432552934654297\n",
            "epoch 18 loss 0.23753149663248369\n",
            "epoch 19 loss 0.21906323085028317\n",
            "epoch 20 loss 0.20713701904300721\n",
            "epoch 21 loss 0.22791918559420493\n",
            "epoch 22 loss 0.20537897133298458\n",
            "epoch 23 loss 0.20976788646751834\n",
            "epoch 24 loss 0.19728194228223256\n",
            "epoch 25 loss 0.21557400706824997\n",
            "epoch 26 loss 0.16888576995341048\n",
            "epoch 27 loss 0.1890002822338213\n",
            "epoch 28 loss 0.20790056818945996\n",
            "epoch 29 loss 0.18659309327842727\n",
            "epoch 30 loss 0.21366153423103593\n",
            "epoch 31 loss 0.18375746563317313\n",
            "epoch 32 loss 0.19759324357484379\n",
            "epoch 33 loss 0.22411791124050656\n",
            "epoch 34 loss 0.22111198902430554\n",
            "epoch 35 loss 0.1782106225183534\n",
            "epoch 36 loss 0.1977246786062155\n",
            "epoch 37 loss 0.1933813807779863\n",
            "epoch 38 loss 0.20827236924801143\n",
            "epoch 39 loss 0.22293430819146096\n",
            "epoch 40 loss 0.19522140822523543\n",
            "epoch 41 loss 0.1946269539034655\n",
            "epoch 42 loss 0.17278945446014404\n",
            "epoch 43 loss 0.19287927464146407\n",
            "epoch 44 loss 0.1907776410601312\n",
            "epoch 45 loss 0.20871546831462653\n",
            "epoch 46 loss 0.2098287620371388\n",
            "epoch 47 loss 0.19232515786443988\n",
            "epoch 48 loss 0.19893933650887302\n",
            "epoch 49 loss 0.26431930894332545\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "daV3W6hDK1Kz"
      },
      "source": [
        "## Inference"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "u02rAgE2K27R",
        "outputId": "da9b6239-51a5-4159-9057-7a33424fa3f3"
      },
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import torchio as tio\n",
        "patch_overlap = 0\n",
        "patch_size = 24, 24, 24\n",
        "target_patch_size = 3\n",
        "\n",
        "#input sampling\n",
        "grid_sampler = tio.inference.GridSampler(subject_list[0], patch_size, patch_overlap)\n",
        "patch_loader = torch.utils.data.DataLoader(grid_sampler, batch_size=4)\n",
        "# target vol sampling\n",
        "grid_sampler_target = tio.inference.GridSampler(subject_list[0], target_patch_size, patch_overlap)\n",
        "aggregator = tio.inference.GridAggregator(grid_sampler_target)\n",
        "target_loader = torch.utils.data.DataLoader(grid_sampler_target, batch_size=4)\n",
        "\n",
        "model.eval()\n",
        "\n",
        "with torch.no_grad():\n",
        "  for patches_batch,target_patches in zip(patch_loader,target_loader):\n",
        "\n",
        "    input_t1 = patches_batch['t1'][tio.DATA]  \n",
        "    input_t2 = patches_batch['t2'][tio.DATA]\n",
        "    input_tensor = torch.cat([input_t1, input_t2], dim=1).float().cuda()\n",
        "\n",
        "    locations = target_patches[tio.LOCATION]\n",
        "    logits = model(input_tensor)\n",
        "    labels = logits.argmax(dim=tio.CHANNELS_DIMENSION, keepdim=True)\n",
        "    outputs = labels\n",
        "    aggregator.add_batch(outputs.type(torch.int32), locations)\n",
        "\n",
        "  print('output tensor shape:',outputs.shape)\n",
        "  output_tensor = aggregator.get_output_tensor()\n",
        "  print(output_tensor.shape)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "output tensor shape: torch.Size([4, 1, 3, 3, 3])\n",
            "torch.Size([1, 144, 192, 256])\n"
          ],
          "name": "stdout"
        }
      ]
    }
  ]
}